#!/usr/bin/env python
# -*- coding: UTF-8 -*-
# @Project : chat_model 
# @File    : bce.py
# @IDE     : PyCharm 
# @Author  :ZH
# @Time    : 2025/1/14 18:23
import asyncio
from typing import List

from apps.embedding_utils.base import EmbeddingUtilBase
from apps.utils import async_http_client


class BcEmbedding(EmbeddingUtilBase):
    def __init__(self, api_key: str = None, *args, **kwargs):
        self.api_key = api_key

    async def embedding(self, base_url, text: str, *args, **kwargs):
        embedding_data = await self.batch_embedding(base_url=base_url, text_list=[text])
        if embedding_data:
            return embedding_data[0]
        else:
            return None

    async def batch_embedding(self, base_url, text_list: List[str], *args, **kwargs):
        embedding_data = await async_http_client.post(
            url=base_url,
            headers={'Content-Type': 'application/json'},
            json=[{'content': value, 'index': index} for index, value in enumerate(text_list)]
        )
        if embedding_data.get('code') == 0:
            embedding_data = sorted(embedding_data['data'], key=lambda x: x['index'])
        else:
            return None
        return [i.get('embedding') for i in embedding_data]

